import torch
import gzip
import numpy as np

HEADER_SIZE = 1024

def pack(t):
    nt = np.array(t.to('cpu'))          # 转化为numpy数组
    buf = gzip.compress(nt.tobytes())
    info = {'dtype':str(nt.dtype), 'shape':nt.shape, 'length':len(buf)}
    return info, buf
    

def unpack(info, buf):
    buf = gzip.decompress(buf)
    nt = np.frombuffer(buf, np.dtype(info["dtype"])).reshape(info["shape"])
    return torch.from_numpy(nt)